import random

import matplotlib.pyplot as plt
from keras.datasets.mnist import load_data
from keras import Sequential, layers, activations, optimizers, losses, Model
import numpy as np
import tensorflow as tf

(x_train, y_train), (x_test, y_test) = load_data()
x_train = x_train.reshape(-1, 28, 28, 1) / 255
x_test = x_test.reshape(-1, 28, 28, 1) / 255

onehot_dim = len(set(y_train))

class AlexNet(Model):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv = Sequential([
            layers.Conv2D(filters=6, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D()
        ])
        self.flat = layers.Flatten()
        self.fc = Sequential([
            layers.Dense(64, activation='relu'),
            layers.Dense(64, activation='relu'),
            layers.Dense(10, activation='softmax')
        ])

    def call(self, inputs):
        x = self.conv(inputs)
        x = self.flat(x)
        x = self.fc(x)
        return x


model = AlexNet()
model.build(input_shape=(None, 28, 28, 1))
model.summary()
model.compile(optimizer=optimizers.Adam(), loss=losses.sparse_categorical_crossentropy, metrics='acc')

log = model.fit(x_train, y_train, epochs=10, batch_size=100, validation_data=(x_test, y_test))

train_acc = log.history['acc']
val_acc = log.history['val_acc']
plt.plot(train_acc, color='r')
plt.plot(val_acc, color='g')
plt.show()

index = random.randint(0, len(x_test) - 1)
predicted_labels = np.argmax(model.predict(x_test[index:index + 1]), axis=-1)[0]
plt.imshow(tf.squeeze(x_test[index]), cmap='gray')
plt.title(f'predict val:{predicted_labels}')
plt.show()
